import os
import random
from collections import defaultdict, Counter

# 设置数据路径
data_dir = r'D:\newtrain\JPEGImages/'  # 替换为你的数据目录
output_dir = r'D:\newtrain/'  # 输出 train.txt 和 val.txt 的目录
image_extension = '.jpg'  # 图片扩展名
label_folder = r'D:\newtrain\labels/'  # 标注文件所在的文件夹名称

# 加载所有图片路径到列表中
image_paths = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith(image_extension)]

# 创建字典来存储每个类别的图像路径
class_images = defaultdict(list)

# 遍历所有图片并读取对应的标注文件
for image_path in image_paths:
    # 获取图片的基本名称（不带扩展名）
    base_name = os.path.splitext(os.path.basename(image_path))[0]

    # 构建标注文件路径
    label_path = os.path.join(data_dir, label_folder, base_name + '.txt')

    if not os.path.exists(label_path):
        print(f"Warning: No label file found for {image_path}")
        continue

    with open(label_path, 'r') as f:
        lines = f.readlines()
        classes_in_image = set(int(line.split()[0]) for line in lines if line.strip())

        for class_id in classes_in_image:
            class_images[class_id].append(image_path)

# 确保每个类别至少有一定比例的数据进入验证集
validation_ratio = 0.2  # 20% 的数据作为验证集
train_images = []
val_images = []

for class_id, images in class_images.items():
    num_val = max(1, int(validation_ratio * len(images)))  # 至少选择一张图片作为验证集
    random.shuffle(images)
    val_images.extend(images[:num_val])
    train_images.extend(images[num_val:])

# 打乱训练集和验证集的顺序
random.shuffle(train_images)
random.shuffle(val_images)

# 写入 train.txt 和 val.txt 文件
with open(os.path.join(output_dir, 'train.txt'), 'w') as f:
    for path in train_images:
        f.write(path + '\n')

with open(os.path.join(output_dir, 'val.txt'), 'w') as f:
    for path in val_images:
        f.write(path + '\n')

print(f"Total images: {len(image_paths)}")
print(f"Training images: {len(train_images)}")
print(f"Validation images: {len(val_images)}")

# 可选：检查每个类别的分布情况
class_distribution_train = Counter()
class_distribution_val = Counter()

for image_path in train_images:
    base_name = os.path.splitext(os.path.basename(image_path))[0]
    label_path = os.path.join(data_dir, label_folder, base_name + '.txt')
    with open(label_path, 'r') as f:
        lines = f.readlines()
        classes_in_image = [int(line.split()[0]) for line in lines if line.strip()]
        class_distribution_train.update(classes_in_image)

for image_path in val_images:
    base_name = os.path.splitext(os.path.basename(image_path))[0]
    label_path = os.path.join(data_dir, label_folder, base_name + '.txt')
    with open(label_path, 'r') as f:
        lines = f.readlines()
        classes_in_image = [int(line.split()[0]) for line in lines if line.strip()]
        class_distribution_val.update(classes_in_image)

print("Training set class distribution:", dict(class_distribution_train))
print("Validation set class distribution:", dict(class_distribution_val))